from tqdm import tqdm
import openai
from config import config as cfg
config = cfg.ai
client = openai.OpenAI(api_key=config.api_key, base_url=config.endpoint)


def openai_send(query: str, message: list, use_web=False, config=config):
    me = message
    if use_web:
        from google_search import do_search
        web_res = do_search(query)
        from sl import jsondump
        if web_res:
            print(web_res)
            web_res = jsondump({i['title']: i['content'] for i in web_res})
            me.append({"role": "user", "content": f"以下是web搜索结果：{web_res}"})
    response = client.chat.completions.create(
        model=config.model,
        messages=me,
        temperature=config.temperature,
        stream=False,
        timeout=60
    )
    return response.choices[0].message.content


translate_prompt = ""
with open("prompt.txt", 'r',encoding='utf-8') as f:
    translate_prompt = f.read()
msg = [
    {"role": "system", "content": translate_prompt}
]


def stacked_push(role, text):
    if len(msg) > 10:
        msg.remove(msg[2])
    msg.append({"role": role, "content": text})


def talk_to_ai(text, message=msg, web=False):
    sents = text.split('\n')
    trans = []
    for sent in tqdm(sents, ncols=60):
        if not sent:
            continue
        try:
            stacked_push("user", sent)
            response = openai_send('', msg, use_web=web)
            if response:
                stacked_push("assistant", response)
                trans.append(response)
        except Exception as e:
            print("Exception in talk_to_ai", e)
            print(sent)
    return '\n'.join(trans)
